import numpy as np
from PIL import Image


class CubeData:
    COLOR_COUNT = 6

    # COLOR_ARRAY = [
    #     [255, 0, 0], [0, 255, 0], [0, 0, 255],
    #     [255, 255, 0], [0, 255, 255], [255, 0, 255],
    # ]
    COLOR_ARRAY = [
        [0, 0, 255],
        [255, 232, 0],
        [0, 255, 0],
        [224, 0, 0],
        [255, 128, 0],
        [255, 255, 255],
    ]

    def __init__(self, randint_matrix, diff_arr, cube_size, border_size, image_path, current_opacity):
        self.randint_matrix_array = [randint_matrix]

        self.diff_arr = diff_arr
        self.diff_arr_record = self.diff_arr.copy()
        self.cube_size: int = cube_size
        self.border_size: int = border_size
        self.is_show_shadow = True

        self.global_diff_count = len(diff_arr)
        self.global_try_diff_count = 0
        self.diff_arr_founded = []

        self.gen_diff_matrix()

        # image_path = "./images/dd.jpg"
        # image_path = "./images/cc.jpeg"
        # image_path = None
        randint_rgb_matrix = self.generate_rgb_matrix(self.randint_matrix_array[0], image_path, current_opacity)
        randint_rgb_matrix_diff = self.generate_rgb_matrix(self.randint_matrix_array[1], image_path, current_opacity)
        self.randint_rgb_matrix_array = [randint_rgb_matrix, randint_rgb_matrix_diff]

    def gen_diff_matrix(self):
        randint_matrix = self.randint_matrix_array[0]
        randint_matrix_diff = randint_matrix.copy()

        for diff_pos in self.diff_arr:
            row, col = diff_pos
            randint_plus = np.random.randint(1, CubeData.COLOR_COUNT - 1)
            randint_matrix_diff[row, col] = (randint_matrix[row, col] + randint_plus) % CubeData.COLOR_COUNT

        self.randint_matrix_array.append(randint_matrix_diff)

    @classmethod
    def gen_data(cls, wall_width, wall_height, cube_size, diff_count, border_size=1, image_path=None, current_opacity=0):
        randint_matrix = np.random.randint(0, cls.COLOR_COUNT, size=(wall_height, wall_width))
        # diff_x_arr = np.random.choice(wall_width, size=diff_count, replace=False)
        # diff_y_arr = np.random.choice(wall_height, size=diff_count, replace=False)

        total_count = wall_width * wall_height
        diff_pos = np.random.choice(total_count, size=diff_count, replace=False)
        diff_arr = []
        for pos in diff_pos:
            row = int(pos / wall_width)
            col = int(pos % wall_width)
            diff_arr.append((row, col))

        # np.vstack((diff_x_arr, diff_y_arr)).T
        return CubeData(randint_matrix, diff_arr, cube_size, border_size, image_path, current_opacity)

    @classmethod
    def generate_rgb_matrix(cls, randint_matrix, image_path=None, current_opacity=0):
        h, w = randint_matrix.shape
        randint_rgb_matrix = np.zeros(shape=(h, w, 3), dtype=np.uint8)
        for row in range(h):
            for col in range(w):
                randint_rgb_matrix[row, col] = cls.COLOR_ARRAY[randint_matrix[row, col]]

        if image_path is None or image_path == "":
            return randint_rgb_matrix

        # 这里和图像进行混合
        image = Image.open(image_path)
        # (383, 725, 3) -> (80,   80, 3)
        np_image = np.array(image)
        img_h, img_w, img_c = np_image.shape
        # 按照比较长的边进行缩放
        scale_ratio_h, scale_ratio_w = h / img_h, w / img_w
        scale_ratio = min(scale_ratio_h, scale_ratio_w)
        new_img_h, new_img_w = round(img_h * scale_ratio), round(img_w * scale_ratio)
        new_image = image.resize((new_img_w, new_img_h), Image.ANTIALIAS)
        new_np_image = np.array(new_image)
        row_start, col_start = round((h - new_img_h) / 2), round((w - new_img_w) / 2)

        org_arr = randint_rgb_matrix[row_start:row_start+new_img_h, col_start:col_start+new_img_w]
        img_arr = new_np_image[:, :]

        # opacity = 0.6
        opacity = current_opacity * 0.1
        # rst_arr = np.round(org_arr * img_arr / 255)
        rst_arr = np.round(opacity * img_arr + (1 - opacity) * org_arr)
        randint_rgb_matrix[row_start:row_start+new_img_h, col_start:col_start+new_img_w] = rst_arr
        return randint_rgb_matrix

    def check_pos(self, c_pos):
        self.global_try_diff_count += 1

        remove_index = -1
        for index, target in enumerate(self.diff_arr_record):
            row, col = target
            x, y = col * self.cube_size, row * self.cube_size
            if x < c_pos.x() < x + self.cube_size and y < c_pos.y() < y + self.cube_size:
                remove_index = index
                break

        rst_target = None
        if remove_index != -1:
            if len(self.diff_arr_record) == 1:
                self.global_diff_count = 0
                self.append_founded_target(self.diff_arr_record[remove_index])
                return

            rst_target = self.diff_arr_record.pop(remove_index)
            self.append_founded_target(rst_target)

        self.global_diff_count = len(self.diff_arr_record)
        return rst_target

    def append_founded_target(self, rst_target):
        for target in self.diff_arr_founded:
            if target[0] == rst_target[0] and target[1] == rst_target[1]:
                return

        self.diff_arr_founded.append(rst_target)


if __name__ == '__main__':
    wall_width = 5
    wall_height = 5
    cube_size = 3
    diff_count = 3
    data = CubeData.gen_data(wall_width, wall_height, cube_size, diff_count, 1)
    print(data.randint_matrix_array)
    print('--------------------------')
    print(diff_count)
